import warnings
warnings.filterwarnings('ignore', 'Optimal rotation is not uniquely or poorly defined for the given sets of vectors.')

from tqdm import tqdm
import numpy as np
from pandas import DataFrame
from scipy.spatial.transform import Rotation as R

from .PDB_dotarray_volume import count_points
from .utils import timer, rnd
from .typing import *


def determine_inner(Coordinates:Points, a_pair:Tuple[str, str], a_pair_xyz:Dict[str, Point], a_pair_info:Dict[str, Contact]) -> Surfaces:
    surface_pair = {}
    atom1 = a_pair[0]
    atom2 = a_pair[1]
    xyz: Dict[str, Point] = {}
    
    # 把原子1放在坐标原点
    for atom in a_pair_xyz:
        xyz[atom] = a_pair_xyz[atom] - a_pair_xyz[atom1]
    
    # 原子2到原子1的向量
    vector: Point = xyz[atom2]
    # 计算这个向量与X轴之间的旋转
    rotation, _ = R.align_vectors([[1, 0, 0]], [vector])
    # 旋转所有点，这样原子2会落到X轴上，可能在正方向也可能在负方向
    centers_array = np.asarray(list(xyz.values()))
    rotated_centers_array = rotation.apply(centers_array)
    # 旋转后的点的坐标
    centers: Dict[str, Point] = {key: value for key, value in zip(xyz.keys(), rotated_centers_array)}
    
    R1: float = a_pair_info[atom1][0]
    R2: float = a_pair_info[atom2][0]
    S: DTYPE = centers[atom2][0] # 两个原子之间的距离（但可能是负值，取决于旋转方向）
    # 解方程得到两球交面在X轴上的坐标，作为阈值
    t: DTYPE = (R1**2 - R2**2 + S**2) / (2*S)
    
    atom1_coor = R1 * Coordinates
    atom2_coor = R2 * Coordinates
    # 每个点代表的体积
    a_surf1 = a_pair_info[atom1][1] / len(Coordinates)
    a_surf2 = a_pair_info[atom2][1] / len(Coordinates)
    # 加上原子中心点的坐标，但只有原子2的X轴需要加，因为原子1在原点
    atom2_coor[:, 0] += S
    
    if S > 0:   # 一对原子中的两个会有不同的半径、表面积
        atom1_contacts = atom1_coor[atom1_coor[:, 0] > t] # 原子1和原子2接触的点阵
        atom2_contacts = atom2_coor[atom2_coor[:, 0] < t] # 原子2和原子1接触的点阵
    elif S < 0: # 如果原子2是在X轴负方向上，那就反过来判断
        atom1_contacts = atom1_coor[atom1_coor[:, 0] < t]
        atom2_contacts = atom2_coor[atom2_coor[:, 0] > t]
    
    # 然后去计算重复的次数。在这一步中，去除了它自己的坐标
    atom1_centers = {(c, a_pair_info[c][0]): centers[c] for c in centers}; del atom1_centers[(atom1, a_pair_info[atom1][0])]
    atom2_centers = {(c, a_pair_info[c][0]): centers[c] for c in centers}; del atom2_centers[(atom2, a_pair_info[atom2][0])]
    atom1_counts = count_points(atom1_contacts, atom1_centers)
    atom2_counts = count_points(atom2_contacts, atom2_centers)
    
    surf1: DTYPE = (1 / atom1_counts).sum()   # 加上原子2自己
    surf2: DTYPE = (1 / atom2_counts).sum()
    
    # 把距离顺便放进去
    dist   = rnd(np.abs(S))
    f_surf = rnd(surf1 * a_surf1)
    r_surf = rnd(surf2 * a_surf2)
    type1  = a_pair_info[atom1][-1]
    type2  = a_pair_info[atom2][-1]
    
    surface_pair[a_pair]       = (f_surf, dist, type1, type2)
    surface_pair[a_pair[::-1]] = (r_surf, dist, type2, type1)
    
    return surface_pair


def pdb_dotarray_surface(ref_fp:Path, contact_df:DataFrame, atom_pairs:dict, disable_print=False) -> Surfaces:
    @timer(disable_print=disable_print)
    def count_surface():
        # 原子中心的坐标
        atoms = contact_df['Name']
        atom_coords = np.asarray([
            contact_df['x'].array, 
            contact_df['y'].array,
            contact_df['z'].array,
        ], dtype=DTYPE).T
        contacts_center: Dict[str, Point] = {i: j for i, j in zip(atoms, atom_coords)}
        contacts_dict: Dict[str, Contact] = contact_df[['Name', 'R', 'Surf', 'Type']].set_index('Name').T.to_dict('list')

        # 导入点的坐标
        Coordinates: Points = np.loadtxt(ref_fp, dtype=DTYPE)
        Pairs_Surface = []
        for a_pair in tqdm(atom_pairs, disable=disable_print):
            # 一个原子对，以及与它们接触的所有原子
            a_atom_pairs = atom_pairs[a_pair]
            a_pair_xyz  = {name: contacts_center[name] for name in a_atom_pairs if name in contacts_center}     # 坐标
            a_pair_info = {name: contacts_dict  [name] for name in a_atom_pairs if name in contacts_dict  }     # 半径、表面积、类型
            
            surface_pair = determine_inner(Coordinates, a_pair, a_pair_xyz, a_pair_info)
            Pairs_Surface.append(surface_pair)
        
        surface = {} # 使用update快速合并字典
        for pair in Pairs_Surface:
            surface.update(pair)  

        return surface
    
    return count_surface()
